from dotenv import load_dotenv
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.llms.tongyi import Tongyi
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough


load_dotenv()

# 1.创建模型
model = Tongyi(model='qwen-plus')

# 测试数据，构建文档
documents = [
    Document(
        page_content='猫是独立的宠物，喜欢安静的环境和舒适的休息空间',
        meta_data={'source': '哺乳动物宠物档案'}
    ),
    Document(
        page_content='兔子性格温顺，适合作为家庭宠物，尤其受孩子喜爱',
        meta_data={'source': '小型哺乳动物指南'}
    ),
    Document(
        page_content='鹦鹉聪明伶俐，能够模仿人类语言，是受欢迎的鸟类宠物',
        meta_data={'source': '鸟类饲养手册'}
    ),
    Document(
        page_content='金鱼易于饲养，适合初学者，能为家庭增添生机',
        meta_data={'source': '观赏鱼养殖指南'}
    ),
    Document(
        page_content='仓鼠活泼好动，适合小空间饲养，但需要定期清理笼子',
        meta_data={'source': '小型宠物护理'}
    ),
    Document(
        page_content='乌龟寿命长，饲养简单，是耐性宠物的代表',
        meta_data={'source': '爬行动物百科'}
    ),
    Document(
        page_content='荷兰猪性格温和，喜欢社交，需要较大的活动空间',
        meta_data={'source': '啮齿类宠物档案'}
    ),
    Document(
        page_content='蜥蜴外形独特，适合喜欢异宠的人群，但需要专业饲养知识',
        meta_data={'source': '爬宠饲养指南'}
    ),
    Document(
        page_content='龙猫毛发柔软，性格胆小，需要安静的环境',
        meta_data={'source': '珍稀宠物介绍'}
    ),
    Document(
        page_content='蛇是冷血动物，饲养需谨慎，适合有经验的宠物主人',
        meta_data={'source': '特殊宠物管理'}
    )
]

# 构建向量空间
# vector_store = Chroma.from_documents(documents, embedding=OpenAIEmbeddings())
# vector_store = Chroma.from_documents(documents, embedding=DashScopeEmbeddings(), persist_directory="./chroma_db")
vector_store = Chroma.from_documents(documents, embedding=DashScopeEmbeddings())

# 相似度查询，返回相似的score，score越低相似度越高
# print(vector_store.similarity_search_with_score(query='咖啡猫'))

# 检索器
retriever = RunnableLambda(vector_store.similarity_search).bind(k=2)
# print(retriever.batch(['加菲猫', '鲨鱼']))


# 提示模板
message = """
使用提供的上下文仅回答这个问题:
{question}
上下文:
{context}
"""

prompt_template = ChatPromptTemplate.from_messages([('human', message)])

chain = {'question': RunnablePassthrough(), 'context': retriever} | prompt_template | model

resp = chain.invoke('请介绍一下蛇')
print(resp)
